import torch
import torch.nn as nn
import torch.optim as optim

N = 1
batch_size = 1
seq_len = 5
input_size = 4
hidden_size = 4

# 文本数据转化为one-hot向量
index2char = ['e', 'h', 'l', 'o']
one_hot_lookup = [[1, 0, 0, 0],
                  [0, 1, 0, 0],
                  [0, 0, 1, 0],
                  [0, 0, 0, 1]]

x_data = [1, 0, 2, 2, 3]  # hello
y_data = [3, 1, 2, 3, 2]  # ohlol

x_one_hot = [one_hot_lookup[x] for x in x_data]
print(x_one_hot)

inputs = torch.Tensor(x_one_hot).view(-1, batch_size, input_size)  # (seq_len, batch_size, input_size)
print(inputs.shape)
print(inputs)
labels = torch.LongTensor(y_data).view(-1, 1)
print(labels.shape)
print(labels)


class Model(nn.Module):
    def __init__(self, batch_size, input_size, hidden_size):
        super(Model, self).__init__()
        self.batch_size = batch_size
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.rnn_cell = nn.RNNCell(input_size=self.input_size, hidden_size=self.hidden_size)

    def forward(self, input, hidden):
        hidden = self.rnn_cell(input, hidden)
        return hidden

    def init_hidden(self):
        return torch.zeros(self.batch_size, self.hidden_size)


net = Model(batch_size, input_size, hidden_size)

criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(net.parameters(), lr=0.1)

for epoch in range(15):
    loss = 0
    optimizer.zero_grad()
    hidden = net.init_hidden()
    print('Predicted string: ', end='')
    for input, label in zip(inputs, labels):
        hidden = net(input, hidden)
        loss += criterion(hidden, label)
        # hidden.max(dim=1) - 在隐藏状态张量的第1维（列维度）上找到最大值，返回最大值和对应的索引
        # idx.item() - 将索引张量转换为Python标量值
        # index2char[idx.item()] - 使用索引从字符映射字典中获取对应的字符
        # print(..., end=' ') - 打印字符并在末尾添加空格而不是换行
        # -------------------------------------------------------------------------------------------------------------
        # hidden.max(dim=1) 会**沿着类别维（dim=1）**取最大值：
        # 返回 (values, indices) 两个张量；values 是最大值本身，indices 是最大值所在的类别下标。
        # 下划线 _ 的意思是丢弃不需要的返回值（这里把最大值本身丢掉，只要下标）。
        _, idx = hidden.max(dim=1)
        print(index2char[idx.item()], end=' ')
    # 在每一批次计算以后，进行反向传播，更新参数
    # 这里batch_size=1，因此loss.backward()只会执行一次
    loss.backward()
    optimizer.step()
    print(f"Epoch: {epoch + 1}, Loss: {loss.item()}")
